import torch
import torch.nn as nn
import torch.nn.functional as F

from .bconv import Bconv

class CFM(nn.Module):
    def __init__(self, in_channels1,in_channels, out_channels):
        super(CFM, self).__init__()
        self.conv1x1 = Bconv(in_channels1, in_channels,1)
        self.conv3x3 = Bconv(in_channels, out_channels, 3, padding=1) 

        self.conv1 = Bconv(in_channels, out_channels, 3, padding=1)
        self.conv2 = Bconv(in_channels, out_channels, 3, padding=1)

        self.conv3 = Bconv(in_channels, out_channels, 3, padding=1)
        self.conv4 = Bconv(in_channels, out_channels, 3, padding=1)

        self.conv5 = Bconv(in_channels, out_channels, 3, padding=1)
        self.conv6 = Bconv(in_channels, out_channels, 3, padding=1)

        self.bn = nn.BatchNorm2d(out_channels)
        self.relu = nn.ReLU(inplace=True)

    def forward(self, x1, x2):
        x2 = F.interpolate(x2, size=x1.size()[2:], mode='bilinear', align_corners=True)
        x2 = self.conv1x1(x2)
        x1 = self.conv3x3(x1)

        x1_1 = self.conv1(x1)
        x2_1 = self.conv2(x2)

        x3 = x2_1 * x1_1

        x1_2 = self.conv3(x3)
        x2_2 = self.conv4(x3)

        x1_3 = x1 + x1_2
        x2_3 = x2 + x2_2

        x1_4 = self.conv5(x1_3)
        x2_4 = self.conv6(x2_3)

        out = x1_4 + x2_4
        out = self.relu(self.bn(out))

        return out